import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint
from functools import partial
from einops import rearrange, reduce

from utils.networkHelper import *

class WeightStandardizedConv2d(nn.Conv2d):
    """
    权重标准化卷积模块
    https://arxiv.org/abs/1903.10520
    这种卷积能与组归一化协同工作，提高模型稳定性
    """
    def forward(self, x):
        eps = 1e-5 if x.dtype == torch.float32 else 1e-3

        weight = self.weight
        mean = reduce(weight, "o ... -> o 1 1 1", "mean")
        var = reduce(weight, "o ... -> o 1 1 1", partial(torch.var, unbiased=False))
        normalized_weight = (weight - mean) * (var + eps).rsqrt()

        return F.conv2d(
            x,
            normalized_weight,
            self.bias,
            self.stride,
            self.padding,
            self.dilation,
            self.groups,
        )

class Block(nn.Module):
    """基本卷积块，包含卷积、归一化和激活"""
    def __init__(self, dim, dim_out, groups=8):
        super().__init__()
        self.proj = WeightStandardizedConv2d(dim, dim_out, 3, padding=1)
        self.norm = nn.GroupNorm(groups, dim_out)
        self.act = nn.SiLU()

    def forward(self, x, scale_shift=None):
        x = self.proj(x)
        x = self.norm(x)

        if exists(scale_shift):
            scale, shift = scale_shift
            x = x * (scale + 1) + shift

        x = self.act(x)
        return x

class ResnetBlock(nn.Module):
    """ResNet风格的块，可选时间嵌入调制"""
    def __init__(self, dim, dim_out, *, time_emb_dim=None, groups=8):
        super().__init__()
        self.mlp = (
            nn.Sequential(nn.SiLU(), nn.Linear(time_emb_dim, dim_out * 2))
            if exists(time_emb_dim)
            else None
        )

        self.block1 = Block(dim, dim_out, groups=groups)
        self.block2 = Block(dim_out, dim_out, groups=groups)
        self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()

    def forward(self, x, time_emb=None):
        scale_shift = None
        if exists(self.mlp) and exists(time_emb):
            time_emb = self.mlp(time_emb)
            time_emb = rearrange(time_emb, "b c -> b c 1 1")
            scale_shift = time_emb.chunk(2, dim=1)

        h = self.block1(x, scale_shift=scale_shift)
        h = self.block2(h)
        return h + self.res_conv(x)

class Attention(nn.Module):
    """多头自注意力机制"""
    def __init__(self, dim, heads=4, dim_head=32):
        super().__init__()
        self.scale = dim_head**-0.5
        self.heads = heads
        hidden_dim = dim_head * heads
        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
        self.to_out = nn.Conv2d(hidden_dim, dim, 1)

    def forward(self, x):
        b, c, h, w = x.shape
        qkv = self.to_qkv(x).chunk(3, dim=1)
        q, k, v = map(
            lambda t: rearrange(t, "b (h c) x y -> b h c (x y)", h=self.heads), qkv
        )
        q = q * self.scale

        sim = torch.einsum("b h d i, b h d j -> b h i j", q, k)
        sim = sim - sim.amax(dim=-1, keepdim=True).detach()
        attn = sim.softmax(dim=-1)

        out = torch.einsum("b h i j, b h d j -> b h i d", attn, v)
        out = rearrange(out, "b h (x y) d -> b (h d) x y", x=h, y=w)
        return self.to_out(out)

class LinearAttention(nn.Module):
    """线性复杂度的注意力实现，避免了O(N²)的计算复杂度"""
    def __init__(self, dim, heads=4, dim_head=32):
        super().__init__()
        self.scale = dim_head**-0.5
        self.heads = heads
        hidden_dim = dim_head * heads
        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)

        self.to_out = nn.Sequential(
            nn.Conv2d(hidden_dim, dim, 1),
            nn.GroupNorm(1, dim)
        )

    def forward(self, x):
        b, c, h, w = x.shape
        qkv = self.to_qkv(x).chunk(3, dim=1)
        q, k, v = map(
            lambda t: rearrange(t, "b (h c) x y -> b h c (x y)", h=self.heads), qkv
        )

        q = q.softmax(dim=-2)
        k = k.softmax(dim=-1)

        q = q * self.scale
        context = torch.einsum("b h d n, b h e n -> b h d e", k, v)

        out = torch.einsum("b h d e, b h d n -> b h e n", context, q)
        out = rearrange(out, "b h c (x y) -> b (h c) x y", h=self.heads, x=h, y=w)
        return self.to_out(out)

class PreNorm(nn.Module):
    """在应用函数前先进行归一化的模块"""
    def __init__(self, dim, fn):
        super().__init__()
        self.fn = fn
        self.norm = nn.GroupNorm(1, dim)

    def forward(self, x):
        x = self.norm(x)
        return self.fn(x)

class UNet(nn.Module):
    """
    UNet模型实现，采用与第二个实现相兼容的接口
    支持梯度检查点功能以节省内存使用
    """
    def __init__(
        self,
        c_in=3,                   # 输入通道数
        c_out=3,                  # 输出通道数
        time_dim=256,             # 时间嵌入维度
        base_channels=64,         # 基础通道数
        device="cuda",            # 设备类型
        attention_heads=4,        # 注意力头数
        attention_dim_head=32,    # 每个注意力头的维度
        dim_mults=(1, 2, 4, 8),   # 通道数倍增系数
        resnet_block_groups=8,    # ResNet块中组归一化的组数
    ):
        super().__init__()
        self.device = device
        self.time_dim = time_dim
        self.use_checkpointing = True  # 启用梯度检查点以优化内存
        
        # 确定维度
        self.channels = c_in
        
        init_dim = base_channels
        self.init_conv = nn.Conv2d(c_in, init_dim, 1, padding=0)  # 初始1x1卷积

        # 计算各层通道数
        dims = [init_dim, *map(lambda m: base_channels * m, dim_mults)]
        in_out = list(zip(dims[:-1], dims[1:]))

        # 使用partial固定ResnetBlock的部分参数
        block_klass = partial(ResnetBlock, groups=resnet_block_groups)

        # 时间嵌入处理
        self.time_mlp = nn.Sequential(
            nn.Linear(time_dim, time_dim),
            nn.GELU(),
            nn.Linear(time_dim, time_dim),
        )

        # 下采样路径
        self.downs = nn.ModuleList([])
        # 上采样路径
        self.ups = nn.ModuleList([])
        num_resolutions = len(in_out)

        # 构建下采样路径
        for ind, (dim_in, dim_out) in enumerate(in_out):
            is_last = ind >= (num_resolutions - 1)

            self.downs.append(
                nn.ModuleList(
                    [
                        block_klass(dim_in, dim_in, time_emb_dim=time_dim),
                        block_klass(dim_in, dim_in, time_emb_dim=time_dim),
                        Residual(PreNorm(dim_in, LinearAttention(dim_in, heads=attention_heads, dim_head=attention_dim_head))),
                        Downsample(dim_in, dim_out)
                        if not is_last
                        else nn.Conv2d(dim_in, dim_out, 3, padding=1),
                    ]
                )
            )

        # 中间层处理
        mid_dim = dims[-1]
        self.mid_block1 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim)
        self.mid_attn = Residual(PreNorm(mid_dim, Attention(mid_dim, heads=attention_heads, dim_head=attention_dim_head)))
        self.mid_block2 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim)

        # 构建上采样路径
        for ind, (dim_in, dim_out) in enumerate(reversed(in_out)):
            is_last = ind == (len(in_out) - 1)

            self.ups.append(
                nn.ModuleList(
                    [
                        block_klass(dim_out + dim_in, dim_out, time_emb_dim=time_dim),
                        block_klass(dim_out + dim_in, dim_out, time_emb_dim=time_dim),
                        Residual(PreNorm(dim_out, LinearAttention(dim_out, heads=attention_heads, dim_head=attention_dim_head))),
                        Upsample(dim_out, dim_in)
                        if not is_last
                        else nn.Conv2d(dim_out, dim_in, 3, padding=1),
                    ]
                )
            )

        # 最终输出层
        self.final_res_block = block_klass(base_channels * 2, base_channels, time_emb_dim=time_dim)
        self.final_conv = nn.Conv2d(base_channels, c_out, 1)

    def pos_encoding(self, t, channels):
        """位置编码函数，将标量时间步编码为高维向量"""
        # 确保t的形状是 [batch_size]
        if t.dim() > 1:
            t = t.squeeze(-1)  # 移除多余的维度
            
        # 创建频率向量
        half_dim = channels // 2
        inv_freq = 1.0 / (
            10000 ** (torch.arange(0, half_dim, device=self.device).float() / half_dim)
        )
        
        # 创建外积，确保维度匹配
        # t形状：[batch_size] -> [batch_size, 1]
        # inv_freq形状：[half_dim] -> [1, half_dim]
        emb = t.unsqueeze(1) * inv_freq.unsqueeze(0)
        
        # 应用正弦和余弦函数
        emb_sin = torch.sin(emb)  # [batch_size, half_dim]
        emb_cos = torch.cos(emb)  # [batch_size, half_dim]
        
        # 拼接正弦和余弦部分
        emb = torch.cat([emb_sin, emb_cos], dim=1)  # [batch_size, channels]
        return emb
        
    def _forward_with_checkpointing(self, module_fn, *args):
        """使用梯度检查点进行前向传播，节省内存"""
        if self.use_checkpointing and self.training and any(arg.requires_grad for arg in args if isinstance(arg, torch.Tensor)):
            return checkpoint(module_fn, *args, use_reentrant=False)
        else:
            return module_fn(*args)

    def _module_forward(self, module, x, *args):
        """模块前向传播的通用函数"""
        return module(x, *args)

    def _attn_forward(self, module, x):
        """注意力模块前向传播的专用函数"""
        return module(x)

    def forward(self, x, t):
        """UNet模型的主前向传播函数"""
        # 时间嵌入处理
        t_emb_input = t.clone().detach()  # 直接使用t，确保不改变原始张量
        t_emb_input = self.pos_encoding(t_emb_input, self.time_dim)
        t = self.time_mlp(t_emb_input)
        
        # 初始卷积层
        x = self.init_conv(x)
        r = x.clone()  # 保存一份用于最终残差连接

        # 特征图的缓存列表，用于跳跃连接
        h = []

        # 下采样路径
        for block1, block2, attn, downsample in self.downs:
            x = self._forward_with_checkpointing(self._module_forward, block1, x, t)
            h.append(x)

            x = self._forward_with_checkpointing(self._module_forward, block2, x, t)
            x = self._forward_with_checkpointing(self._attn_forward, attn, x)
            h.append(x)

            x = self._forward_with_checkpointing(self._module_forward, downsample, x)

        # 中间层
        x = self._forward_with_checkpointing(self._module_forward, self.mid_block1, x, t)
        x = self._forward_with_checkpointing(self._attn_forward, self.mid_attn, x)
        x = self._forward_with_checkpointing(self._module_forward, self.mid_block2, x, t)

        # 上采样路径
        for block1, block2, attn, upsample in self.ups:
            x = torch.cat((x, h.pop()), dim=1)  # 添加跳跃连接
            x = self._forward_with_checkpointing(self._module_forward, block1, x, t)

            x = torch.cat((x, h.pop()), dim=1)  # 添加跳跃连接
            x = self._forward_with_checkpointing(self._module_forward, block2, x, t)
            x = self._forward_with_checkpointing(self._attn_forward, attn, x)

            x = self._forward_with_checkpointing(self._module_forward, upsample, x)

        # 最终处理
        x = torch.cat((x, r), dim=1)  # 最后的残差连接
        x = self._forward_with_checkpointing(self._module_forward, self.final_res_block, x, t)
        
        return self.final_conv(x)

def test_unet():
    """测试UNet模型的功能"""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")
    
    batch_size = 2
    channels_in = 3
    channels_out = 3
    height, width = 256, 256
    time_dim = 256
    
    x = torch.randn(batch_size, channels_in, height, width).to(device)
    t_scalar = torch.randint(0, 1000, (batch_size,), device=device).long()
    
    print(f"输入图像形状: {x.shape}")
    print(f"时间步形状: {t_scalar.shape}")

    # 创建模型，使用较小的参数用于测试
    model = UNet(
        c_in=channels_in, 
        c_out=channels_out, 
        time_dim=time_dim, 
        base_channels=32, 
        device=device, 
        attention_heads=2, 
        attention_dim_head=16
    ).to(device)
    
    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"模型总可训练参数数量: {total_params:,}")
    
    try:
        # 测试前向传播
        output = model(x, t_scalar)
        print(f"输出形状: {output.shape}")
        
        assert output.shape == (batch_size, channels_out, height, width), \
            f"期望输出形状 {(batch_size, channels_out, height, width)}，实际得到 {output.shape}"
        
        # 测试梯度反向传播
        if device.type != 'mps':
            loss = output.sum()
            loss.backward()
            print("梯度反向传播测试成功!")

        print("UNet前向传播测试成功!")
        
    except Exception as e:
        print(f"UNet前向传播测试失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    test_unet()